package com.example.yckjbigdataflink.source;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;

import java.util.Arrays;
import java.util.List;

public class HBaseScanSource extends RichParallelSourceFunction<Result> {

    private volatile boolean isRunning = true;
    private Connection connection;
    private Table table;

    @Override
    public void open(Configuration parameters) throws Exception {
        org.apache.hadoop.conf.Configuration config = HBaseConfiguration.create();
        config.set("hbase.zookeeper.quorum", "zk1,zk2,zk3");
        config.set("zookeeper.znode.parent", "/hbase");
        connection = ConnectionFactory.createConnection(config);
        table = connection.getTable(TableName.valueOf("my_hbase_table"));
    }

    @Override
    public void run(SourceContext<Result> ctx) throws Exception {
        int totalTasks = getRuntimeContext().getNumberOfParallelSubtasks();
        int taskIndex = getRuntimeContext().getIndexOfThisSubtask();

        // 切分 rowkey 范围，假设前缀是数字或字母
        List<String> rowKeyPrefixes = Arrays.asList("00", "01", "...", "99"); // 分片前缀
        int step = rowKeyPrefixes.size() / totalTasks;

        List<String> assignedPrefixes = rowKeyPrefixes.subList(
                taskIndex * step,
                (taskIndex == totalTasks - 1) ? rowKeyPrefixes.size() : (taskIndex + 1) * step
        );

        for (String prefix : assignedPrefixes) {
            Scan scan = new Scan();
            scan.withStartRow(Bytes.toBytes(prefix));
            scan.withStopRow(Bytes.toBytes(prefix + "~")); // ASCII 排序 trick
            scan.setCaching(5000);
            scan.setCacheBlocks(false); // 禁用 block 缓存，避免冲击 regionserver

            ResultScanner scanner = table.getScanner(scan);
            for (Result result : scanner) {
                if (!isRunning) break;
                ctx.collect(result);
            }
            scanner.close();
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
    }

    @Override
    public void close() throws Exception {
        if (table != null) table.close();
        if (connection != null) connection.close();
    }
}

